package heap02

func NewHeap(cap int, compare func(interface{}, interface{}) bool) MyHeap {
	return MyHeap{
		cap: cap,
		compare: compare,
		nodes: make([]interface{}, cap),
		size: 0,
	}
}

type MyHeap struct {
	cap     int
	size    int
	nodes   []interface{}
	compare func(interface{}, interface{}) bool
}

func (mH *MyHeap) Size() int {
	return mH.size
}

func (mH *MyHeap)Add(node interface{}) {
	if mH.size == mH.cap {
		panic(any("heap is full"))
	}
	mH.nodes[mH.size] = node
	mH.Insert(mH.size)
	mH.size++
}

func (mH *MyHeap)Insert(index int) {
	for mH.compare(mH.nodes[index], mH.nodes[(index-1) / 2]) {
		mH.Swap(index, (index-1) / 2)
		index = (index-1) / 2
	}
}

func (mH *MyHeap) Poll() interface{} {
	if mH.size == 0 {
		panic(any("heap is empty"))
	}
	ans := mH.nodes[0]
	mH.size--
	mH.Swap(0, mH.size)
	mH.heapify(0)
	return ans
}

func (mH *MyHeap) heapify(index int) {
	cur := index
	next := (cur<<1) | 1
	for next < mH.size {
		if next+1 < mH.size && !mH.compare(mH.nodes[next], mH.nodes[next+1]) {
			next = next +1
		}
		if !mH.compare(mH.nodes[cur], mH.nodes[next]) {
			mH.Swap(cur, next)
		} else {
			break
		}
		cur = next
		next = (cur<<1) | 1
	}
}

func (mH *MyHeap)Peek() interface{} {
	if mH.size == 0 {
		return nil
	}
	return mH.nodes[0]
}

func (mH *MyHeap)Swap(index1, index2 int) {
	mH.nodes[index1], mH.nodes[index2] = mH.nodes[index2], mH.nodes[index1]
}